#!/usr/bin/env bash

# Eases deployment of lwaftr inside a virtual machine.
#
# The scripts enables several commands:
#
# * lwaftrctl start|stop snabbnfv   Starts or stops snabbnfv.
# * lwaftrctl start|stop vm         Starts or stops virtual machine
#
# A configuration file named 'lwaftrctl.conf' must exist in the current
# directory. This file contains all the variable definitions needed to run the
# available commands. See 'lwaftrctl.conf.example'.
#
# The usual workflow to run the script would be the following:
#
# * lwaftrctl start snabbnfv.
#       Brings up NICs in host that will be used by virtual machine.
#
# * lwaftrctl start vm.
#       Brings up VM. After this step it should be possible to log into
#       the VM: telnet localhost 5001.
#
# The command 'lwaftrctl all start', run all the steps above.

QEMU=qemu-system-x86_64
LWAFTR_DELAY=${LWAFTR_DELAY:-"20"}

tmux_session=""
snabb_n=0

function run_telnet {
    (echo "$2"; sleep ${3:-2}) \
        | telnet localhost $1 2>&1
}

function tmux_launch {
    local id=$1
    local command="$2 2>&1 | tee $3"
    tmux_session="$id-$$"
    tmux new-session -d -n "$id" -s $tmux_session "$command"
}

function qemu_start_vm {
    echo "Started QEMU"
    cd $VM_DIR
    tmux_launch \
        "qemu0" \
        "sudo numactl -m ${NUMA_NODE} taskset -c ${QEMU_CORE} \
            $QEMU $QEMU_ARGS \
            -kernel ${BZ_IMAGE} \
                -append \"earlyprintk root=/dev/vda rw console=ttyS1 ${KERNEL_PARAMS} ip=${VM_IP} hugepages=256\" \
            -M pc -smp 1 -cpu host -m ${MEM} -enable-kvm \
            -fsdev local,security_model=passthrough,id=fsdev0,path=${SHARED_LOCATION} \
            -device virtio-9p-pci,id=fs0,fsdev=fsdev0,mount_tag=share \
            -numa node,memdev=mem -object memory-backend-file,id=mem,size=${MEM},mem-path=${HUGEPAGES_FS},share=on \
            -netdev type=vhost-user,id=net0,chardev=char0 -chardev socket,id=char0,path=${SNABBNFV_SOCK1},server \
                -device virtio-net-pci,netdev=net0,addr=0x8,mac=${SNABBNFV_MAC1} \
            -netdev type=vhost-user,id=net1,chardev=char1 -chardev socket,id=char1,path=${SNABBNFV_SOCK2},server \
                -device virtio-net-pci,netdev=net1,addr=0x9,mac=${SNABBNFV_MAC2} \
            -serial telnet:localhost:${VM_PORT},server,nowait \
            -serial stdio \
            -display none \
            -drive if=virtio,format=${VM_FORMAT},file=${VM_IMAGE} " \
        "qemu0.log"
}

function start_vm {
    qemu_start_vm
}

function stop_vm {
    echo "Stopping VM"
    run_telnet ${VM_PORT} "sudo poweroff" >/dev/null
    pid=$(pidof $QEMU)
    sudo kill ${pid} 2>/dev/null
    echo "Done"
}

function restart_vm {
    stop_vm
    start_vm
}

function start_snabbnfv {
    start_snabbnfv_process "snabbnfv-1" $SNABBNFV_CORE1 $SNABBNFV_PCI1 $SNABBNFV_CONF1 $SNABBNFV_SOCK1
    start_snabbnfv_process "snabbnfv-2" $SNABBNFV_CORE2 $SNABBNFV_PCI2 $SNABBNFV_CONF2 $SNABBNFV_SOCK2
}

function start_snabbnfv_process {
    local id=$1
    local core=$2
    local pci=$3
    local conf=$4
    local sock=$5

    # Check configuration file exists.
    if [[ ! -f ${conf} ]]; then
        echo "WARNING: File '${conf}' does not exist."
    fi

    # Run snabbnfv inside tmux session.
    tmux_launch "$id" \
        "cd ${SNABB_PATH}/src; sudo numactl -m ${NUMA_NODE} taskset -c ${core} ./snabb snabbnfv traffic -b ${pci} ${conf} ${sock}" \
        $(snabb_log)
    status=$(tmux ls | grep -c "$id")
    sleep 0.5

    # Check exit status.
    if [[ ${status} -eq 0 ]]; then
        echo "Start of snabbnfv failed: "
        echo -e "\tsudo numactl -m ${NUMA_NODE} taskset -c ${core} ./snabb snabbnfv traffic -b ${pci} ${conf} ${sock}"
        exit 1
    fi

    echo "Started snabbnfv on core ${core} (tmux: '$id')"
    echo -e "\t{PCI: ${pci}, conf: ${conf}, sock: ${sock}}"
}

function snabb_log {
    echo "snabb${snabb_n}.log"
}

function stop_snabbnfv {
    echo "Stopping snabbnfv"
    kill_all "snabbnfv traffic"
    remove_file $SNABBNFV_SOCK1
    remove_file $SNABBNFV_SOCK2
    echo "Done"
}

function kill_all {
    local name=$1
    pids=`ps aux | grep "$name" | awk '{ print $2 }'`
    for pid in ${pids[@]}; do
        sudo kill $pid 2>/dev/null
    done
    sleep 1
}

function remove_file {
    sudo rm -f $1
}

function restart_snabbnfv {
    stop_snabbnfv
    start_snabbnfv
}

function start_lwaftr {
    echo "Started lwaftr"
    run_telnet ${VM_PORT} "sudo ~/start-lwaftr.sh" >/dev/null
}

function stop_lwaftr {
    echo "Stopping lwaftr"
    run_telnet ${VM_PORT} "sudo ~/stop-lwaftr.sh" >/dev/null
    echo "Done"
}

function restart_lwaftr {
    stop_lwaftr
    start_lwaftr
}

# Actions.

function start_command {
    COMMAND=$1
    case $COMMAND in
        "all")
            start_snabbnfv
            start_vm
            echo "Waiting ${LWAFTR_DELAY} seconds"
            sleep $LWAFTR_DELAY
            start_lwaftr
            ;;
        "snabbnfv")
            start_snabbnfv
            ;;
        "vm")
            start_vm
            echo "Connect via 'telnet localhost ${VM_PORT}'"
            ;;
        "lwaftr")
            start_lwaftr
            ;;
        *)
            bad_usage
            ;;
    esac
}

function stop_command {
    COMMAND=$1
    case $COMMAND in
        "all")
            stop_lwaftr
            stop_vm
            stop_snabbnfv
            ;;
        "snabbnfv")
            stop_snabbnfv
            ;;
        "vm")
            stop_vm
            ;;
        "lwaftr")
            stop_lwaftr
            ;;
        *)
            bad_usage
            ;;
    esac
}

function restart_command {
    COMMAND=$1
    case $COMMAND in
        "all")
            restart_snabbnfv
            restart_vm
            ;;
        "snabbnfv")
            restart_snabbnfv
            ;;
        "vm")
            restart_vm
            ;;
        "lwaftr")
            restart_lwaftr
            ;;
        *)
            bad_usage
            ;;
    esac
}

# Main

function usage {
    local exit_code=$1
    echo "Usage: lwaftrctl -f lwaftrctl.conf [start|stop|restart] [all|snabbnfv|vm|lwaftr]"
    exit $exit_code
}

function bad_usage {
    usage -1
}

function help {
    usage 0
}

AFTRCTL_CONF=lwaftrctl.conf
ARGS=()
while [[ $# > 0 ]]; do
    key="$1"

    case $key in
        -f|--file)
            AFTRCTL_CONF="$2"
            shift
            ;;
        *)
            ARGS+=($key)
            ;;
    esac
    shift
done

if [ ${#ARGS[@]} -ne 2 ]; then
    help
fi

if [ ! -f "$AFTRCTL_CONF" ]; then
    echo "Could not find lwaftrctl.conf in current directory. See lwaftrctl.conf.example."
    exit -1
fi

source $AFTRCTL_CONF

PROGRAM_NAME=$0
ACTION=${ARGS[0]}
COMMAND=${ARGS[1]}

case $ACTION in
    "start")
        start_command $COMMAND;
        ;;
    "stop")
        stop_command $COMMAND;
        ;;
    "restart")
        restart_command $COMMAND;
        ;;
    *)
        echo "Unknown action: $ACTION"
        exit -1
        ;;
esac

exit 0
